#ifndef AER_DOFEFFECT_HPP_
#define AER_DOFEFFECT_HPP_

#include "aerPostEffect.hpp"


namespace aer {
  
class DoFEffect : public PostEffect
{
  private:
    struct Pass
    {
      Program program;
      RenderTexture rt;
      Texture color;
    };
    
    struct Input
    {
      Texture *pColor;
      Texture *pDepth;
      Vector2 *pLinearParams;
    };
    
    
  private:
    Pass m_coeffPass; 
    Pass m_downsamplePass;
    Pass m_blurPass;
    Pass m_compoPass; 
    
    Input m_input;
    
    Vector2 m_texelSize;
    float m_focusPlaneDistance;
		float m_focusRange;
    
  
  public:
    DoFEffect()
    {
      m_input.pColor = NULL;
      m_input.pDepth = NULL;
      m_input.pLinearParams = NULL;
      
      // Inputs
      addInput( "Color", m_input.pColor);
      addInput( "Depth", m_input.pDepth);
      addInput( "LinearParams", m_input.pLinearParams);
      
      // Output
      addOutput( "Color", m_compoPass.color);
      
      
      m_focusPlaneDistance = 500.0f;
      m_focusRange = 500.0f;
    }
    
    ~DoFEffect()
    {
      //
    }
    
    
    void init()
    {
      // DoF COEFF
      const char *coeffFS = AER_STRINGIFY(
        \x23 version 330\n
        
        in vec2 vTexCoord;
        out vec4 fragColor;
        
        uniform sampler2D uDepthTex;
        uniform vec2 uLinearParams;
        uniform float uFocusPlaneDist;
        uniform float uNearRange;
        uniform float uFarRange;


        void main()
        {
          float viewDepth = texture( uDepthTex, vTexCoord ).r;
          viewDepth = uLinearParams.y / (viewDepth - uLinearParams.x);
          
          float d = viewDepth - uFocusPlaneDist;
          
          if (d < 0.0f) {
            d = abs( d / uNearRange );
          } else {
            d = d / uFarRange;
          }
          
          d = clamp( d, 0.0f, 1.0f);
          fragColor = vec4( d, 0.0f, 0.0f, 1.0f);
        }
      );
      
      m_coeffPass.program.generate();
        m_coeffPass.program.addShaderSrc( Program::VERTEX_SHADER, ScreenQuad::GetVertexShader());
        m_coeffPass.program.addShaderSrc( Program::FRAGMENT_SHADER, coeffFS);
      m_coeffPass.program.link();
      
      
      // DOWNSAMPLE
      const char *downsampleFS = AER_STRINGIFY(
        \x23 version 330\n
        
        in vec2 vTexCoord;
        
        out vec4 fragColor;        
        uniform sampler2D uInputTex;
        
        void main() { fragColor = texture( uInputTex, vTexCoord); }
      );
      
      m_downsamplePass.program.generate();
        m_downsamplePass.program.addShaderSrc( Program::VERTEX_SHADER, ScreenQuad::GetVertexShader());
        m_downsamplePass.program.addShaderSrc( Program::FRAGMENT_SHADER, downsampleFS);
      m_downsamplePass.program.link();
      
      
      // BLUR
      const char *blurFS = AER_STRINGIFY(
        \x23 version 330\n
        \x23 define PI 3.14159f\n
        
        in vec2 vTexCoord;
        out vec4 fragColor;
        
        uniform sampler2D uInputTex;
        uniform sampler2D uCoeffTex;
        uniform vec2 uBlurDir;   // wether (texelWidth, 0.0f) OR (0.0f, texelHeight)
        uniform int uKernelRadius = 1;


        void main()
        {
          float blurCoeff = texture( uCoeffTex, vTexCoord).r;
          blurCoeff = 2.0f * blurCoeff * blurCoeff;
          float invBlurCoeff = 1.0f / blurCoeff;
          
          vec4 blurColor = vec4(0.0f);
          float sum = 0.0f;
          for (int r = -uKernelRadius; r <= +uKernelRadius; ++r)
          {
            vec2 uv = vTexCoord + r * uBlurDir;
            
            float d = (uBlurDir.x + uBlurDir.y) * r; // only make sense with orthogonal direction
            float w = inversesqrt( PI * blurCoeff ) * exp( - d*d*invBlurCoeff );
            
            blurColor += texture( uInputTex, uv) * w;
            sum += w;
          }
          
          fragColor = blurColor / sum;
        }
      );
      
      m_blurPass.program.generate();
        m_blurPass.program.addShaderSrc( Program::VERTEX_SHADER, ScreenQuad::GetVertexShader());
        m_blurPass.program.addShaderSrc( Program::FRAGMENT_SHADER, blurFS);
      m_blurPass.program.link();
      
      
      // COMPOSITION
      const char *compoFS = AER_STRINGIFY(
        \x23 version 330\n
        
        in vec2 vTexCoord;
        out vec4 fragColor;
        
        uniform sampler2D uInputTex;
        uniform sampler2D uBlurCoeffTex;
        uniform sampler2D uDownsampleTex;
        
        void main()
        {
          float blurCoeff = texture( uBlurCoeffTex, vTexCoord).r;
          
          vec4 inputColor = texture( uInputTex, vTexCoord);
          vec4 downsampleColor = texture( uDownsampleTex, vTexCoord);
          
          fragColor = mix( inputColor, downsampleColor, blurCoeff);
          //fragColor.a = inputColor.a;//
        }
      );
      
      m_compoPass.program.generate();
        m_compoPass.program.addShaderSrc( Program::VERTEX_SHADER, ScreenQuad::GetVertexShader());
        m_compoPass.program.addShaderSrc( Program::FRAGMENT_SHADER, compoFS);
      m_compoPass.program.link();
      
      
      m_coeffPass.color.generate();
      m_downsamplePass.color.generate();
      m_blurPass.color.generate();
      m_compoPass.color.generate();
      
      m_coeffPass.color.bind();
        m_coeffPass.color.setFilters( Texture::LINEAR, Texture::LINEAR);
        m_coeffPass.color.setWraps( Texture::CLAMP_TO_EDGE, Texture::CLAMP_TO_EDGE);
      m_downsamplePass.color.bind();
        m_downsamplePass.color.setFilters( Texture::LINEAR, Texture::LINEAR);
        m_downsamplePass.color.setWraps( Texture::CLAMP_TO_EDGE, Texture::CLAMP_TO_EDGE);
      m_blurPass.color.bind();
        m_blurPass.color.setFilters( Texture::LINEAR, Texture::LINEAR);
        m_blurPass.color.setWraps( Texture::CLAMP_TO_EDGE, Texture::CLAMP_TO_EDGE);
      m_compoPass.color.bind();
        m_compoPass.color.setFilters( Texture::NEAREST, Texture::NEAREST);
        m_compoPass.color.setWraps( Texture::CLAMP_TO_EDGE, Texture::CLAMP_TO_EDGE);
      Texture::Unbind( Texture::TEXTURE_2D, 0u);
      
      
      m_coeffPass.rt.generate();
      m_downsamplePass.rt.generate();
      m_blurPass.rt.generate();
      m_compoPass.rt.generate();
    }
    
    void resize(const U32 width, const U32 height)
    {
      U32 subWidth = width >> 1;
      U32 subHeight = height >> 1;
      
      m_coeffPass.color.bind();
      m_coeffPass.color.resize( width, height, GL_R16F);
      
      m_downsamplePass.color.bind();
      m_downsamplePass.color.resize( subWidth, subHeight, m_input.pColor->getInternalFormat());
      
      m_blurPass.color.bind();
      m_blurPass.color.resize( subWidth, subHeight, m_downsamplePass.color.getInternalFormat());
      
      m_compoPass.color.bind();
      m_compoPass.color.resize( width, height, m_blurPass.color.getInternalFormat());
      
      Texture::Unbind( Texture::TEXTURE_2D, 0u);
      
      
      m_coeffPass.rt.bind();
      AER_CHECK( m_coeffPass.rt.setup( &m_coeffPass.color ) );
      m_downsamplePass.rt.bind();
      AER_CHECK( m_downsamplePass.rt.setup( &m_downsamplePass.color ) );
      m_compoPass.rt.bind();
      AER_CHECK( m_compoPass.rt.setup( &m_compoPass.color ) );
      RenderTarget::unbind();
      
      
      m_texelSize = Vector2( 1.0f/subWidth, 1.0f/subHeight );
    }
    
    void apply()
    {
      AER_ASSERT( "Color input not specified" && m_input.pColor != NULL );
      AER_ASSERT( "Depth input not specified" && m_input.pDepth != NULL );
      AER_ASSERT( "LinearParam input not specified" && m_input.pLinearParams != NULL );
      
      
      glViewport( 0, 0, m_input.pColor->getWidth(), m_input.pColor->getHeight());
      
      /// 1) Blur coeffs
      m_coeffPass.rt.begin();
      m_coeffPass.program.bind();
        m_coeffPass.program.setUniform( "uDepthTex", 0);
        m_coeffPass.program.setUniform( "uLinearParams", *m_input.pLinearParams);
        m_coeffPass.program.setUniform( "uFocusPlaneDist", m_focusPlaneDistance);
        m_coeffPass.program.setUniform( "uNearRange", 2.0f * m_focusRange / 3.0f);
        m_coeffPass.program.setUniform( "uFarRange", m_focusRange);
        m_input.pDepth->bind( 0u );
          ScreenQuad::Draw();                
        m_input.pDepth->unbind();        
      //m_coeffPass.program.unbind();
      //m_coeffPass.rt.end();
      
      
      glViewport( 0, 0, m_downsamplePass.color.getWidth(), m_downsamplePass.color.getHeight());
      
      /// 2) Downsample
      m_downsamplePass.rt.begin();
      m_downsamplePass.program.bind();
        m_downsamplePass.program.setUniform( "uDepthTex", 0);
        m_input.pColor->bind( 0u );
          ScreenQuad::Draw();
        m_input.pColor->unbind();
      //m_downsamplePass.program.unbind();
      //m_downsamplePass.rt.end();
      
      
      /// 3) Blur (X + Y) the downsample
      m_blurPass.rt.begin();
      m_blurPass.program.bind();
            
      m_blurPass.program.setUniform( "uKernelSize",  4);//
      m_blurPass.program.setUniform( "uInputTex",  0);
      m_blurPass.program.setUniform( "uCoeffTex",  1);
      m_coeffPass.color.bind( 1u );
      
      // X-blur
      m_blurPass.program.setUniform( "uBlurDir", Vector2( m_texelSize.x, 0.0f) );
      m_downsamplePass.color.bind( 0u );                      // Tex Input
      AER_CHECK( m_blurPass.rt.setup( &m_blurPass.color ) );  // Tex Output
        ScreenQuad::Draw();
        
      // Y-blur
      m_blurPass.program.setUniform( "uBlurDir", Vector2( 0.0f, m_texelSize.y) );
      m_blurPass.color.bind( 0u );                                  // Tex Input
      AER_CHECK( m_blurPass.rt.setup( &m_downsamplePass.color ) );  // Tex Output
        ScreenQuad::Draw();
      
      Texture::Unbind( Texture::TEXTURE_2D, 1u);
      Texture::Unbind( Texture::TEXTURE_2D, 0u);
      
      m_blurPass.program.unbind();
      m_blurPass.rt.end();
      
      
      glViewport( 0, 0, m_compoPass.color.getWidth(), m_compoPass.color.getHeight());
      
      /// 4) Compo
      if (getTarget() == TEXTURE) { m_compoPass.rt.begin(); }
      m_compoPass.program.bind();
      m_compoPass.program.setUniform( "uInputTex",      0);
      m_compoPass.program.setUniform( "uBlurCoeffTex",  1);
      m_compoPass.program.setUniform( "uDownsampleTex", 2);
      
      m_input.pColor->bind( 0u );
      m_coeffPass.color.bind( 1u );
      m_downsamplePass.color.bind( 2u );
      
        ScreenQuad::Draw();
      
      Texture::Unbind( Texture::TEXTURE_2D, 2u);
      Texture::Unbind( Texture::TEXTURE_2D, 1u);
      Texture::Unbind( Texture::TEXTURE_2D, 0u);
      
      m_compoPass.program.unbind();
      if (getTarget() == TEXTURE) { m_compoPass.rt.end(); }
      
      
      CHECKGLERROR();
    }
    
    
    void setFocusPlaneDistance(const float focusPlaneDist) { m_focusPlaneDistance = focusPlaneDist; }
    void setFocusRange(const float focusRange) { m_focusRange = focusRange; }
    
    float getFocusPlaneDistance() const { return m_focusPlaneDistance; }
    float getFocusRange() const { return m_focusRange; }
};
  
}

#endif //AER_DOFEFFECT_HPP_
